#!/usr/bin/python
# Author: Jimmy Saw
# Date: 06/04/2011
# Renames fasta ids based on a given prefix
# Usage example: python rename_fasta_ids.py seqs.fasta GLBK > seqs.newids.fasta
# Currently goes from 0001 to 9999. Change zfill number to higher if needed.

import sys
import re
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord

fasta_file = sys.argv[1]
out_prefix = sys.argv[2]
moltype = sys.argv[3]
seqs = SeqIO.parse(fasta_file, "fasta")

for i, seq in enumerate(seqs):
    new_id = out_prefix + "_" + str(i+1).zfill(4)
    new_seq = seq.seq
    tmprec = SeqRecord(new_seq, id=new_id, 
        description="")
    if moltype == "protein":
        print tmprec.format("fasta").strip('*\n')
    else:
        print tmprec.format("fasta").strip('\n')
